#include <stdio.h>
#include <string.h>
#include <fcntl.h>
#include <sys/ioctl.h>
#include <linux/kvm.h>
#include <unistd.h>
#include <sys/mman.h>
#include <stddef.h>
#include <stdlib.h>

// 调整宏定义条件，确保正确性
#ifndef KVM_EXIT_ARM_SMC
#define KVM_EXIT_ARM_SMC 23
#endif

#ifndef KVM_EXIT_ARM_HVC
#define KVM_EXIT_ARM_HVC 22
#endif

#define KVM_DEV  "/dev/kvm"
#define MEM_SIZE  0x1000
#define PHY_ADDR  0x80000
 
#define AARCH64_CORE_REG(x)        (KVM_REG_ARM64 | KVM_REG_SIZE_U64 | KVM_REG_ARM_CORE | KVM_REG_ARM_CORE_REG(x))

static void panic(const char *msg) {
    perror(msg);
    exit(EXIT_FAILURE);
}
 

static int vcpu_set_reg(int vcpufd, __u64 reg_id, __u64 value)
{
    struct kvm_one_reg reg_struct;
    memset(&reg_struct, 0, sizeof(struct kvm_one_reg));
    reg_struct.id = reg_id;
    __u64 *value_ptr = malloc(sizeof(__u64));
    *value_ptr = value;
    reg_struct.addr = (__u64)value_ptr;
 
    //set the vcpu register
    int ret = ioctl(vcpufd, KVM_SET_ONE_REG, &reg_struct);
    if (ret < 0) {
        perror("KVM_SET_ONE_REG");
        free(value_ptr);
        return -1;
    }
    free(value_ptr);
    return 0;

}

int main(void)
{
    struct kvm_one_reg reg;
    struct kvm_vcpu_init init; //using init the vcpu type
    struct kvm_vcpu_init preferred;
    int ret;
 
    __u64 guest_entry = PHY_ADDR;
    __u64 guest_pstate;
 
    int kvmfd = open(KVM_DEV, O_RDWR);
    if(kvmfd < 0) {
        panic("open kvm failed\n");
    }
 
    //ioctl(kvmfd, KVM_GET_API_VERSION, NULL);
    //1. create vm and get the vm fd handler
    int vmfd =  ioctl(kvmfd, KVM_CREATE_VM, 0);
    if(vmfd < 0) {
        panic("create vm failed\n");
    }
 
    //2. create vcpu
    int vcpufd = ioctl(vmfd, KVM_CREATE_VCPU, 0);
    if(vcpufd < 0) {
        panic("create vcpu failed\n");
    }
 
    //3. arm64 type vcpu type init
    //sample code can check the qemu/target/arm/kvm64.c
    memset(&init, 0, sizeof(init));
    //init.target = KVM_ARM_TARGET_GENERIC_V8; //here set KVM_ARM_TARGET_CORTEX_A57 meet failed
    init.target = -1;
    if (init.target == -1) {
        ret = ioctl(vmfd, KVM_ARM_PREFERRED_TARGET, &preferred);
        if(!ret) {
            init.target = preferred.target; //KVM_ARM_TARGET_GENERIC_V8
            printf("preferred vcpu type %d\n", init.target);
        } else {
            printf("get preferred vcpu type failed\n");
        }
    }
    
    // 使用最基础的vCPU配置，不启用任何特殊特性
    init.features[0] = 0;
    
    ret = ioctl(vcpufd, KVM_ARM_VCPU_INIT, &init);
 
    if(ret  < 0) {
        perror("init vcpu type failed");
        panic("init vcpu type failed\n");
    } else {
        printf("vCPU initialized successfully with features: 0x%x\n", init.features[0]);
    }
 
    //4. get vcpu resouce map size and get vcpu  kvm_run status  
    int mmap_size = ioctl(kvmfd, KVM_GET_VCPU_MMAP_SIZE, NULL);
 
    if(mmap_size  < 0) {
        panic("get vcpu mmap size failed\n");
    }
 
    struct kvm_run *run = mmap(NULL, mmap_size, PROT_READ | PROT_WRITE, MAP_SHARED, vcpufd, 0);
    if(run == MAP_FAILED) {
        panic("mmap vcpu run failed\n");
    }
 
    //5. load the vm running program to buffer 'ram'
    unsigned char *ram = mmap(NULL, MEM_SIZE, PROT_READ | PROT_WRITE, MAP_SHARED | MAP_ANONYMOUS, -1, 0);
 
    int kfd = open("test.bin", O_RDONLY);
    if(kfd < 0) {
        panic("open test.bin failed\n");
    }
    read(kfd, ram, MEM_SIZE);
         
    struct kvm_userspace_memory_region mem = {
        .slot = 0,
        .flags = 0,
        .guest_phys_addr = PHY_ADDR,
        .memory_size = MEM_SIZE,
        .userspace_addr = (unsigned long)ram,
    };
 
    //6. set the vm userspace program ram to vm fd handler
    ret = ioctl(vmfd, KVM_SET_USER_MEMORY_REGION, &mem);
 
    if(ret < 0) {
        panic("set user memory region failed\n");
    }
 
    //7. change the vcpu register info, arm64 need change the pc value 
    // arm64 not support KVM_SET_REGS, and KVM_SET_SREGS only support at x86 ppc arch
    // arm64 using the KVM_SET_ONE_REG
 
    //sample code from qemu/target/arm/kvm64.c
    //reg.id = AARCH64_CORE_REG(regs.pc);
    //reg.addr = (__u64)&guest_entry;
    //ret = ioctl(vcpufd, KVM_SET_ONE_REG, &reg);

    //if(ret  < 0) {
    //    panic("change arm64 pc reg failed err %d\n");
    //} else {
    //    printf("set pc addr success\n");
    //}

    vcpu_set_reg(vcpufd, AARCH64_CORE_REG(regs.pc), guest_entry);

    vcpu_set_reg(vcpufd, AARCH64_CORE_REG(regs.sp), PHY_ADDR + MEM_SIZE);

    vcpu_set_reg(vcpufd, AARCH64_CORE_REG(spsr[2]), 0x5);  // EL1h模式

    //vcpu_set_reg(vcpufd, AARCH64_CORE_REG(regs.pstate), 0x5); // set pstate 0x5, EL1h mode
 
    //7. run the vcpu and get the vcpu result
    printf("Starting VM execution...\n");
    int run_count = 0;
    while(1) {
        ret = ioctl(vcpufd, KVM_RUN, NULL);
        if (ret == -1)
        {
            perror("KVM_RUN failed");
            panic("exit unknow\n");
        }
        
        run_count++;
        if (run_count > 100) {  // 增加循环次数限制
            printf("Too many KVM_RUN iterations, exiting\n");
            return -1;
        }
 
        switch(run->exit_reason)
        {
            //when guest program meet io access error will trigger KVM_EXIT_MMIO event, 
            //using the event get guest program output
            case KVM_EXIT_MMIO:
                if (run->mmio.phys_addr == 0x996 && run->mmio.is_write && run->mmio.len == 1) {
                    printf("%c", run->mmio.data[0]);
                    if (run->mmio.data[0] == '\n') {
                        printf("\nVM execution completed successfully via MMIO.\n");
                        return 0;
                    }
                }
                break;
            
            case KVM_EXIT_HYPERCALL:
                printf("KVM_EXIT_HYPERCALL\n");
                return 0;

            case KVM_EXIT_FAIL_ENTRY:
                printf("Entry error: hw_entry_failure = 0x%llx\n", run->fail_entry.hardware_entry_failure_reason);
                return -1;
                
            // 处理SMC指令退出
            case KVM_EXIT_ARM_SMC:
                printf("Received SMC call - exiting VM\n");
                return 0;
                
            // 修改HVC指令处理方式，尝试继续执行而不是退出
            case KVM_EXIT_ARM_HVC:
                printf("Received HVC call - exiting VM\n");
                return 0;

            default:
                printf("Unhandled exit reason: %d\n", run->exit_reason);
                // 尝试继续执行而不是立即退出
                continue;

        }
    }
    return 0;
}